{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0f8f7dfb",
   "metadata": {},
   "source": [
    "# Optimize fetching data from Neo4j with Apache Arrow\n",
    "## High-performance data retrieval from Neo4j with Apache Arrow\n",
    "\n",
    "The year is 2022, and graph machine learning is one of the rising trends in data analytics. While Neo4j has a Graph Data Science library that supports multiple graph algorithms and machine learning workflows, sometimes you want to export data from Neo4j and run it through your favorite machine learning frameworks like PyTorch or TensorFlow. In that scenario, you want to be able to export data from Neo4j in a fast and scalable way. But, unfortunately, using the Neo4j Python driver is not the most efficient way of retrieving data. However, no need to worry, Dave Voutila has got your back. In the past couple of months, he has been developing an Apache Arrow plugin for Neo4j.\n",
    "\n",
    "https://github.com/neo4j-field/neo4j-arrow\n",
    "\n",
    "\n",
    "The goal of the Neo4j Arrow project is to expose data available in Neo4j via high-performance Arrow Flight APIs. You can retrieve data via Cypher queries or even fetch information from the GDS in-memory graphs. I am not familiar with the underlying Apache Arrow infrastructure, so I won't try to explain how it works. Just know that it is blazing fast.\n",
    "## Preparing the environment\n",
    "We will use a subset of the Pokec dataset to demonstrate the data retrieval of Neo4j Arrow. The dataset contains 1 million nodes and 10 million relationships. You can download the Neo4j [database dump at this link](https://drive.google.com/file/d/176-Vmdn2fqy4KLPygl-yjJLiIhMXzKOL/view). To instantiate the database from the dump, open the Neo4j Desktop application and copy the dump under the Files section. Next, click on the three dots in the right-hand side of the file, and select Create new DBMS from dump.\n",
    "\n",
    "It is crucial that you select a database version of 4.4.0 or later, as otherwise, the Arrow plugin might not work. Currently, the Neo4j Arrow project is still in the early stages, so not all database versions are supported. I have prepared the latest build of the [Arrow plugin for download](https://drive.google.com/file/d/1Xi0jwqFiJx_aKtZZ-Rd4XZrpMrJke6u5/view), or you can build it locally yourself. To activate the plugin, simply copy it in the database plugins folder.\n",
    "To follow this blog post, you will also need to install the Graph Data Science library.\n",
    "\n",
    "## Retrieve data via Cypher query\n",
    "You can now start the database instance and learn how to retrieve data via Arrow plugin. I have prepared a Jupyter notebook that you can use to follow the examples.\n",
    "First, we will use the official Neo4j Python driver to retrieve all the nodes and their properties from the database. As mentioned, there are around a million nodes in the database dump I prepared. The Cypher statement we will use is pretty straightforward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4268519d",
   "metadata": {},
   "outputs": [],
   "source": [
    "cypher = \"\"\"\n",
    "MATCH (u:User)\n",
    "RETURN u.id AS id, u.gender AS gender, u.age AS age\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64b8bf58",
   "metadata": {},
   "source": [
    "Now, we can go ahead and execute the Cypher statement. We will also measure the execution time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "849ade8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from neo4j import GraphDatabase\n",
    "import time\n",
    "\n",
    "host = 'bolt://localhost:7687'\n",
    "user = 'neo4j'\n",
    "password = 'letmein'\n",
    "driver = GraphDatabase.driver(host,auth=(user, password))\n",
    "\n",
    "def execute_query(query):\n",
    "    with driver.session() as session:\n",
    "        data = list(session.run(query))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b89e6d32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neo4j driver took 32.6 seconds to fetch the data\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "data = execute_query(cypher)\n",
    "delta = round(time.time() - start, 1)\n",
    "print(f'Neo4j driver took {delta} seconds to fetch the data')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "409e8f1f",
   "metadata": {},
   "source": [
    "The Neo4j Python driver fetches a million rows in 30 seconds, which is not that bad. However, when dealing with tens or hundreds of millions of rows, it doesn't scale well enough. Here is where the Neo4j Arrow plugin comes in handy.\n",
    "The nice benefit of the Arrow plugin is that it is language-independent, and you could use it to retrieve data from several scripting languages. At the moment, there is a Python wrapper object available, but you could develop the Arrow client in other languages as well.\n",
    "The code to fetch data using a Cypher statement via the Python Arrow wrapper is pretty straightforward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b03ad428",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neo4j arrow took 2.8 seconds to fetch the data\n"
     ]
    }
   ],
   "source": [
    "import neo4j_arrow as na\n",
    "client = na.Neo4jArrow(user, password, ('localhost', 9999))\n",
    "\n",
    "start = time.time()\n",
    "ticket = client.cypher(cypher)\n",
    "table = client.stream(ticket).read_all()\n",
    "delta = round(time.time() - start, 1)\n",
    "\n",
    "print(f'Neo4j arrow took {delta} seconds to fetch the data')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fc3704b",
   "metadata": {},
   "source": [
    "It took 2.8 seconds to retrieve a million rows from the Neo4j database on my laptop. While this is only a 10x improvement, the Arrow plugin is more scalable and will outperform the Python driver up to 450x on larger datasets.\n",
    "Since data gurus have a soft spot for Pandas, you can easily convert the table object to a Pandas dataframe using the `to_pandas()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "38eadc26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>gender</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>26</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>19</td>\n",
       "      <td>1</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>0</td>\n",
       "      <td>17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>36</td>\n",
       "      <td>1</td>\n",
       "      <td>28</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>39</td>\n",
       "      <td>1</td>\n",
       "      <td>18</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   id  gender  age\n",
       "0   1       1   26\n",
       "1  19       1   21\n",
       "2  21       0   17\n",
       "3  36       1   28\n",
       "4  39       1   18"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = table.to_pandas()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efa1b2dd",
   "metadata": {},
   "source": [
    "## Retrieve data from GDS in-memory graph\n",
    "What is remarkable about the Neo4j Arrow project is that it also allows high-performance reads from the Graph Data Science library in-memory graphs. You could perform feature engineering using any of the graph or embedding algorithms, or you could simply evaluate the algorithm's result in Python.\n",
    "We will construct a GDS in-memory graph and execute the Weakly Connected Components algorithm to demonstrate how to retrieve data efficiently from the GDS in-memory graph. Let's start by constructing the in-memory graph called pokec that contains nodes with a label User and FRIEND relationships."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a01b35c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<Record nodeProjection={'User': {'label': 'User', 'properties': {}}} relationshipProjection={'FRIEND': {'orientation': 'NATURAL', 'aggregation': 'DEFAULT', 'type': 'FRIEND', 'properties': {}}} graphName='pokec' nodeCount=1099121 relationshipCount=10794057 createMillis=2679>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "execute_query(\"\"\"\n",
    "CALL gds.graph.create('pokec', 'User', 'FRIEND')\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "827d94c8",
   "metadata": {},
   "source": [
    "Now, we will execute the mutate mode of the Weakly Connected Components algorithm to store the results back to the in-memory graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1000092a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<Record mutateMillis=1 nodePropertiesWritten=1099121 componentCount=921 componentDistribution={'p99': 5, 'min': 2, 'max': 1097079, 'mean': 1193.3984799131379, 'p90': 3, 'p50': 2, 'p999': 30, 'p95': 3, 'p75': 2} postProcessingMillis=99 createMillis=0 computeMillis=765 configuration={'seedProperty': None, 'consecutiveIds': False, 'threshold': 0.0, 'relationshipWeightProperty': None, 'nodeLabels': ['*'], 'sudo': False, 'relationshipTypes': ['*'], 'mutateProperty': 'wcc', 'username': None, 'concurrency': 4}>]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "execute_query(\"\"\"\n",
    "CALL gds.wcc.mutate('pokec', {mutateProperty:'wcc'})\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cd4f6df",
   "metadata": {},
   "source": [
    "With the prepared GDS in-memory graph, we can use the Arrow plugin to extract the algorithm results from Neo4j. Thanks to Dave Voutila and the Python wrapper he prepared, this is a pretty simple process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aacfb373",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">> Reading the result of our GDS job...\n",
      ">> Took 1.85s to consume stream into a PyArrow table.\n"
     ]
    }
   ],
   "source": [
    "# Submit our GDS job to retrieve some node properties from a graph projection\n",
    "print('>> Reading the result of our GDS job...''')\n",
    "start = time.time()\n",
    "ticket = client.gds_nodes('pokec', properties=['wcc'])\n",
    "# Retrieve and consume the stream into a PyArrow Table\n",
    "table = client.stream(ticket).read_all()\n",
    "delta = round(time.time() - start, 2)\n",
    "print(f'>> Took {delta:,}s to consume stream into a PyArrow table.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f84a756a",
   "metadata": {},
   "source": [
    "We retrieved the WCC results from the projected in-memory graph in 1.8 seconds. You could now use the WCC algorithm results to efficiently split the train-test data, use it as a feature in a downstream ML workflow, or simply evaluate and plot the results.\n",
    "Of course, the popular way of interacting with data in Python is using the Pandas dataframe. We can easily transform the retrieved data into a dataframe with the `to_pandas()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6ec771ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   _node_id_ _labels_  wcc\n",
      "0       1365   [User]    0\n",
      "1       1378   [User]    0\n",
      "2       1386   [User]    0\n",
      "3       1393   [User]    0\n",
      "4       1400   [User]    0\n",
      "921\n"
     ]
    }
   ],
   "source": [
    "wcc_df = table.to_pandas()\n",
    "print(wcc_df.head())\n",
    "print(wcc_df['wcc'].nunique())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab8d7e0c",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "The Neo4j Arrow project is still in the early stages, so please test it out and report any bugs or feature requests directly to the [repository](https://github.com/neo4j-field/neo4j-arrow), and if you like it, please give it a star. I must say I am very excited about this project and how it can help graph data scientists to efficiently retrieve data from Neo4j, which makes it more enjoyable to implement various machine learning pipelines around Neo4j."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab602d07",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
